import os.path
import time
import cv2
import folder_paths

__VERSION__ = "v1.3.4"

class SimpleWebpSaverNode:
    CATEGORY = "AmorForge"
    # RETURN_TYPES = ("IMAGE",)
    # RETURN_NAMES = ("image_output_name",)

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                # "High Quality": (["enable", "disable"]),
                "quality": ("INT", {
                    "default": 60,
                    "min": 30, #Minimum value
                    "max": 90, #Maximum value
                    "display": "slider"
                }),
                "save_to_folder": ("STRING", {
                    "multiline": False,
                    "default": "webp"
                }),
                "reverse_rgb": ("BOOLEAN", {
                    "default": True
                }),
            },
        }

    def _get_latest_counter(self, save_to_folder):
        """ Get the latest counter from the files in the save_to_folder. """
        files = os.listdir(save_to_folder)
        # counters = [int(f.split(".")[0]) for f in files if f.endswith(".webp")]
        list_counters = []
        for f in files:
            # if f.endswith(".webp"):
            base_name, ext = os.path.splitext(f)
            if ext == ".webp":
                try:
                    counter = int(base_name.rsplit("-", 1)[1])
                    list_counters.append(counter)
                except Exception:
                    continue
        return max(list_counters) if list_counters else 0

    def _get_unique_fname(self, idx):
        # 结合当前时间戳，确保文件名唯一
        timestamp = time.time()
        str_timestamp = f"{timestamp:.4f}"
        return f"{str_timestamp}-{idx}.webp"

    def run(self, image, quality, save_to_folder, reverse_rgb=False):
        """ type(image) is: torch.Tensor """
        output_dir = folder_paths.get_output_directory()
        save_dir = os.path.join(output_dir, save_to_folder)
        os.makedirs(save_dir, exist_ok=True)

        images_np = image.cpu().numpy() * 255.

        results = []
        # counter = self._get_latest_counter(save_dir)
        for idx, img_np in enumerate(images_np, start=1):
            img_uint8 = img_np.astype("uint8")

            # 如果需要反转RGB通道为BGR
            if reverse_rgb:
                img_uint8 = img_uint8[:, :, ::-1]  # 反转通道顺序 RGB -> BGR

            if 1:
                save_rpath = self._get_unique_fname(idx)
            else:
                counter += 1
                save_rpath = f"AmorComfyUI-{counter}.webp"
            save_fpath = os.path.join(save_dir, save_rpath)
            cv2.imwrite(save_fpath, img_uint8, [cv2.IMWRITE_WEBP_QUALITY, quality])

            results.append({
                'filename': save_rpath,
                'subfolder': save_to_folder,
                'type': 'output'})
        return {'ui': {'images': results}}
